package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/go-ping/ping"
	"github.com/pion/stun"
	"gitlab.com/NebulousLabs/go-upnp"
)

// Check if an IP is private IPv4
func isPrivateIP(ip net.IP) bool {
	if ip4 := ip.To4(); ip4 != nil {
		switch {
		case ip4[0] == 10:
			return true
		case ip4[0] == 172 && ip4[1]&0xf0 == 16:
			return true
		case ip4[0] == 192 && ip4[1] == 168:
			return true
		}
	}
	return false
}

// Returns the first IPv4 of iface or dials outbound address to guess local IP/interface
func getLocalIPAndInterface(ifaceName, externalDial string) (net.IP, *net.Interface, error) {
	if ifaceName != "" {
		iface, err := net.InterfaceByName(ifaceName)
		if err != nil {
			return nil, nil, err
		}
		for _, a := range mustGetAddrs(iface) {
			if ip := a.To4(); ip != nil {
				return ip, iface, nil
			}
		}
		return nil, nil, errors.New("no IPv4 address on specified interface")
	}
	conn, err := net.Dial("udp", externalDial)
	if err != nil {
		return nil, nil, err
	}
	defer conn.Close()
	localAddr := conn.LocalAddr().(*net.UDPAddr)
	ip := localAddr.IP
	ifaces, _ := net.Interfaces()
	for _, iface := range ifaces {
		for _, a := range mustGetAddrs(&iface) {
			if a.Equal(ip) {
				return ip, &iface, nil
			}
		}
	}
	return ip, nil, nil
}

func mustGetAddrs(iface *net.Interface) []net.IP {
	addrs, _ := iface.Addrs()
	ips := make([]net.IP, 0, len(addrs))
	for _, a := range addrs {
		if ipnet, ok := a.(*net.IPNet); ok {
			ips = append(ips, ipnet.IP)
		}
	}
	return ips
}

// Send a basic STUN Binding request (with optional CHANGE-REQUEST attrs)
func sendBinding(server string, timeout time.Duration, changeIP, changePort bool, localIP net.IP) (*stun.XORMappedAddress, int, error) {
	if !strings.Contains(server, ":") {
		server += ":3478"
	}
	laddr := &net.UDPAddr{IP: localIP, Port: 0}
	conn, err := net.ListenUDP("udp", laddr)
	if err != nil {
		return nil, 0, err
	}
	defer conn.Close()
	msg := stun.MustBuild(stun.TransactionID, stun.BindingRequest)
	if changeIP || changePort {
		var flag byte
		if changeIP { flag |= 0x04 }
		if changePort { flag |= 0x02 }
		msg.Add(stun.AttrChangeRequest, []byte{0, 0, 0, flag})
	}
	msg.NewTransactionID()
	conn.SetDeadline(time.Now().Add(timeout))
	host, portStr, _ := net.SplitHostPort(server)
	port := 3478
	fmt.Sscanf(portStr, "%d", &port)
	if _, err := conn.WriteTo(msg.Raw, &net.UDPAddr{IP: net.ParseIP(host), Port: port}); err != nil {
		return nil, 0, err
	}
	buf := make([]byte, 1500)
	n, err := conn.Read(buf)
	if err != nil {
		return nil, 0, err
	}
	var res stun.Message
	res.Raw = buf[:n]
	if err := res.Decode(); err != nil {
		return nil, 0, err
	}
	var xor stun.XORMappedAddress
	if err := xor.GetFrom(&res); err != nil {
		return nil, 0, err
	}
	localPort := conn.LocalAddr().(*net.UDPAddr).Port
	return &xor, localPort, nil
}

// Filter alive STUN servers by basic bind+timeout/concurrency
func filterServers(servers []string, timeout time.Duration, concurrency int, localIP net.IP) []string {
	var mu sync.Mutex
	live := make([]string, 0, len(servers))
	sem := make(chan struct{}, concurrency)
	var wg sync.WaitGroup
	for _, srv := range servers {
		wg.Add(1)
		srv := srv
		go func() {
			defer wg.Done()
			sem <- struct{}{}
			if _, _, err := sendBinding(srv, timeout, false, false, localIP); err == nil {
				mu.Lock()
				live = append(live, srv)
				mu.Unlock()
			}
			<-sem
		}()
	}
	wg.Wait()
	return live
}

// Multi-server NAT type classification (RFC3489 stacking)
func testPair(s1, s2 string, timeout time.Duration, localIP net.IP) string {
	if _, _, err := sendBinding(s1, timeout, false, false, localIP); err != nil {
		return "UDP blocked or unreachable"
	}
	if _, _, err := sendBinding(s1, timeout, true, true, localIP); err == nil {
		return "Full Cone NAT"
	}
	if _, _, err := sendBinding(s2, timeout, false, false, localIP); err == nil {
		if _, _, err2 := sendBinding(s1, timeout, false, true, localIP); err2 == nil {
			return "Restricted Cone NAT"
		}
		return "Port-Restricted Cone NAT"
	}
	return "Symmetric NAT"
}

// Default STUN list
var defaultSTUNServers = []string{
	"222.137.252.18",
	"103.69.128.47",
	"49.0.239.69",
	"223.244.91.19",
	"103.219.34.70",
	"103.219.34.72",
	"8.218.148.237",
	"125.70.231.96",
	"218.22.52.100",
	"120.193.211.252",
	"39.153.135.73",
	"59.44.218.162",
	"119.117.142.4",
	"60.14.52.186",
	"111.2.168.150",
	"221.7.195.52",
	"111.53.11.104",
	"122.9.67.165",
	"116.63.110.231",
	"1.192.136.170",
	"111.198.39.10",
	"180.101.234.103",
	"223.111.68.6",
	"223.111.68.4",
}

// Print eth0 MAC/IPs if available
func printEth0IPAndMAC() {
	iface, err := net.InterfaceByName("eth0")
	if err != nil {
		fmt.Println("eth0 interface not found:", err)
		return
	}
	fmt.Printf("eth0 MAC: %s\n", iface.HardwareAddr.String())
	addrs, err := iface.Addrs()
	if err != nil {
		fmt.Println("Failed to get eth0 addresses:", err)
		return
	}
	for _, addr := range addrs {
		fmt.Printf("eth0 IP: %s\n", addr.String())
	}
}

// Check IPv6 connectivity by ping6
func checkIPv6Connectivity() {
	ipv6Target := "240c::6666"
	fmt.Printf("Checking IPv6 connectivity by pinging %s ...\n", ipv6Target)
	pinger, err := ping.NewPinger(ipv6Target)
	if err != nil {
		fmt.Println("failed to create pinger:", err)
		return
	}
	pinger.Count = 3
	pinger.SetPrivileged(true)
	pinger.SetNetwork("udp6")
	err = pinger.Run()
	if err != nil {
		fmt.Println("failed to run ping:", err)
		return
	}
	stats := pinger.Statistics()
	if stats.PacketLoss < 100 {
		fmt.Println("IPv6 is available! (loss:", stats.PacketLoss, "%)")
	} else {
		fmt.Println("IPv6 is NOT available.")
	}
}

// Try UPnP IGD (WAN) gateway and external IP
func testUPNPGateway() {
	fmt.Print("Checking for UPnP IGD gateway... ")
	ch := make(chan *upnp.IGD, 1)
	errch := make(chan error, 1)
	go func() {
		d, err := upnp.Discover()
		if err != nil {
			errch <- err
			return
		}
		ch <- d
	}()
	select {
	case d := <-ch:
		ip, err := d.ExternalIP()
		if err != nil {
			fmt.Println("Found gateway, but failed to query external IP")
			return
		}
		fmt.Printf("Found! Gateway: %s, external IP: %s\n", d.Location(), ip)
	case err := <-errch:
		fmt.Println("Not found or not supported (", err, ")")
	case <-time.After(5 * time.Second):
		fmt.Println("UPnP discovery timed out after 5 seconds.")
	}
}

func readServersFromFile(filePath string) ([]string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()
	return parseServerListFromReader(file), nil
}

func parseServerListFromReader(r io.Reader) []string {
	servers := []string{}
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && !strings.HasPrefix(line, "#") {
			servers = append(servers, line)
		}
	}
	return servers
}

// ------- MAIN -------

func main() {
	filePath   := flag.String("file", "", "path to STUN server list")
	timeout    := flag.Duration("timeout", 3*time.Second, "request timeout duration")
	conc       := flag.Int("concurrency", 20, "concurrent probes for filtering")
	maxPairs   := flag.Int("maxPairs", 11, "maximum server pairs to test")
	ifaceName  := flag.String("iface", "", "network interface name (optional)")
	externalDial := flag.String("dial", "223.6.6.6:53", "external address to dial for local IP detection")
	flag.Parse()

	printEth0IPAndMAC()
	checkIPv6Connectivity()
	testUPNPGateway()

	// NAT detect base
	localIP, iface, err := getLocalIPAndInterface(*ifaceName, *externalDial)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: failed to get local IP: %v\n", err)
		os.Exit(1)
	}
	if iface != nil {
		fmt.Printf("Using interface %s with IP %s\n", iface.Name, localIP)
	} else {
		fmt.Printf("Using IP %s\n", localIP)
	}

	// Load STUN servers
	var servers []string
	if *filePath != "" {
		servers, err = readServersFromFile(*filePath)
		if err != nil || len(servers) == 0 {
			fmt.Fprintf(os.Stderr, "Warning: can't read servers from %s, fallback to built-in list.\n", *filePath)
			servers = nil
		} else {
			fmt.Printf("Loaded %d servers from local file: %s\n", len(servers), *filePath)
		}
	}
	if len(servers) == 0 {
		servers = defaultSTUNServers
		fmt.Println("No valid STUN IPs file provided, using built-in default list and checking for live servers...")
	}

	// NAT type detection preflow:
	// 1. Direct public IP (no NAT at all)
	if !isPrivateIP(localIP) {
		fmt.Printf("Local IP %s appears public, verifying with STUN...\n", localIP)
		mapped, lp, err := sendBinding(servers[0], *timeout, false, false, localIP)
		if err == nil && mapped.IP.Equal(localIP) && mapped.Port == lp {
			fmt.Println("Detected direct public IP (no NAT)")
			os.Exit(0)
		}
		// Otherwise, continues below (possible SNAT)
	}
	// 2. Early "UPnP/DMZ/public mapping" (compatible with cloud, static port mapping, etc)
	mapped, lp, err := sendBinding(servers[0], *timeout, false, false, localIP)
	if err == nil && !isPrivateIP(mapped.IP) && mapped.Port == lp {
		fmt.Printf("Detected public port mapping (UPnP/DMZ/cloud static NAT): public IP %s, port %d\n", mapped.IP.String(), mapped.Port)
		os.Exit(0)
	}

	// 3. Multi-pair NAT classification (RFC3489 logic)
	live := filterServers(servers, *timeout, *conc, localIP)
	if len(live) < 2 {
		fmt.Fprintln(os.Stderr, "Error: not enough active STUN servers")
		os.Exit(1)
	}
	fmt.Printf("Active STUN servers: %d\n", len(live))
	rand.Seed(time.Now().UnixNano())
	counts := make(map[string]int)
	pairsTested := 0
	usedPairs := make(map[[2]int]bool)
	for pairsTested < *maxPairs {
		i := rand.Intn(len(live))
		j := rand.Intn(len(live))
		if i == j {
			continue
		}
		a, b := i, j
		if b < a {
			a, b = b, a
		}
		if usedPairs[[2]int{a, b}] {
			continue
		}
		usedPairs[[2]int{a, b}] = true
		res := testPair(live[i], live[j], *timeout, localIP)
		fmt.Printf("%s <-> %s => %s\n", live[i], live[j], res)
		counts[res]++
		pairsTested++
	}

	var finalType string
	var maxCnt int
	for t, c := range counts {
		if c > maxCnt {
			maxCnt = c
			finalType = t
		}
	}
	fmt.Printf("\nFinal NAT type after %d random pair tests: %s\n", pairsTested, finalType)
}


